# llm_engine/build_index.py
import os
import sys
import django

from llm_engine.config import CHROMA_DB_PATH

from llama_index.core import Settings
from llama_index.core import Document, VectorStoreIndex, StorageContext
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.embeddings.zhipuai import ZhipuAIEmbedding
from llama_index.llms.zhipuai import ZhipuAI
from chromadb.utils.embedding_functions import DefaultEmbeddingFunction
from chromadb import Documents, EmbeddingFunction, Embeddings
import chromadb 

from llm_engine.config import API_KEY, EMBEDDING_MODEL_NAME


def setup_django_env():
    """
    一个独立的工具函数，设置Django环境。
    它假设Django项目'the_site'与'llm_engine'是兄弟目录。
    """
    # 获取当前脚本所在目录 (llm_engine)
    llm_engine_dir = os.path.dirname(os.path.abspath(__file__))
    # 获取父目录 (项目的根目录)
    project_root = os.path.dirname(llm_engine_dir)
    # Django项目的路径
    django_project_path = os.path.join(project_root, 'the_site')
    
    if django_project_path not in sys.path:
        sys.path.append(django_project_path)
    
    os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'the_site.settings')
    try:
        django.setup()
        print("Django environment setup successfully!")
    except ImportError as e:
        raise ImportError(
            "Couldn't import Django. Are you sure it's installed and "
            "available on your PYTHONPATH environment variable? Did you "
            "forget to activate a virtual environment?"
        ) from e


class ZhipuAIEmbeddingFunction(EmbeddingFunction):
    """适配智谱AI Embeddings的EmbeddingFunction"""
    def __call__(self, input: Documents) -> Embeddings:
        # 调用智谱AI的嵌入模型API
        # 示例代码，具体API调用方式请参考智谱AI官方文档/SDK
        embed_model = ZhipuAIEmbedding(
            api_key=API_KEY,
            model=EMBEDDING_MODEL_NAME
        )
        embeddings = []
        for text in input:
            # 提取1024维向量
            emb = embed_model.get_text_embedding(text)
            embeddings.append(emb)
        return embeddings

def main():
    """主执行函数"""
    print("Setting up Django environment to access models...")
    setup_django_env()
    
    # 只有在setup成功后才能导入模型
    from project.models import Project
    from users.models import Profile
    
    def get_department_for_owner(owner):
        profile = Profile.objects.get(user=owner)
        return profile.department.name
    print("Models imported. Starting to build knowledge base...")
    
    # === 1) 读取 Django 项目数据，构造伪文档 ===
    documents = []
    for project in Project.objects.all().order_by("id"):
        text = f"项目名称: {project.name}\n项目描述: {project.description}"
        metadata={"department": get_department_for_owner(project.owner), "project_id": str(project.id)} # 核心在这里, 插入metadata
        doc = Document(
                text=text,
                metadata=metadata,
                doc_id=f"project_{project.id}" # 防止多次embedding重复.
            )
        documents.append(doc)
    if not documents:
        print("No Project records found.")
        return
    
    # === 2) 初始化 ZhipuAI Embedding ===
    embed_model = ZhipuAIEmbedding(
        api_key=API_KEY,
        model=EMBEDDING_MODEL_NAME
    )
    
    # 实例化你的嵌入函数
    zhipu_ef = ZhipuAIEmbeddingFunction()
    
    Settings.embed_model = embed_model
    
    # === 3) 准备 Chroma 持久化客户端 ===
    chroma_client = chromadb.PersistentClient(path=CHROMA_DB_PATH) # 这里用Client不会落盘
    chroma_collection = chroma_client.get_or_create_collection("project_collection", embedding_function=zhipu_ef)
    print(type(chroma_collection)) # 调试为啥不落盘.
    print(chroma_collection._embedding_function)
    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
    
    # -----------------------------
    # 4) 手动添加文档到 collection（保证落盘）
    # -----------------------------
    chroma_collection.add(
        documents=[doc.text for doc in documents],
        metadatas=[{
            "doc_id": doc.doc_id,
            "department": doc.metadata["department"],
        } for doc in documents],
        ids=[doc.doc_id for doc in documents],
    )
    
    # -----------------------------
    # 5) 构建 ChromaVectorStore 并创建 VectorStoreIndex
    # -----------------------------
    vector_store = ChromaVectorStore(chroma_collection=chroma_collection)

    index = VectorStoreIndex.from_documents(
        documents,
        vector_store=vector_store,
        show_progress=True,
    )
    
    # -----------------------------
    # 6) 查询或调试 collection
    # -----------------------------
    all_docs = chroma_collection.get(include=["documents", "metadatas", "embeddings"])
    for doc, meta, eid in zip(all_docs['documents'], all_docs['metadatas'], all_docs['ids']):
        print(f"ID: {eid}, Metadata: {meta}, Document: {doc}")

    # === 7)（可选）简单查询示例 ===
    # 下面示例展示如何在命令里做一次 quick sanity check
    try:
        llm = ZhipuAI(
                api_key=API_KEY,
                model="glm-4-plus",   # 智谱通用模型
                temperature=0.3
            )
        
        query_engine = index.as_query_engine(
            similarity_top_k=3,
            llm=llm
        )
        resp = query_engine.query("这个库里有哪些和‘水利部’相关的项目？")
        print("Sample query result:")
        print(str(resp))
    except Exception as e:
        print(f"Query test skipped: {e}")
    
    
    


if __name__ == "__main__":
    main()